// Copyright 2022 jeff.li. and/or its affiliates.

#pragma once

#include <unordered_map>
#include <unordered_set>
#include <vector>

#include <tbir/ir/hlo_builtin.h>
#include <tbir/ir/stmt_functor.h>

namespace tbir::ir {

    class FuseContAnyGetSetItemOptimizer : public StmtExprMutator {
    public:
        BaseFunc run(const BaseFunc &func) {
            return this->VisitExpr(func);
        }

        HLOExpr VisitExpr_(const CallNode *op) override {
            if (op->op.same_as(builtin::object___setitem__())) {
                MXCHECK(op->args.size() == 3) << "internal error";
                runtime::Array<BaseExpr> keys;
                keys.push_back(op->args[1]);
                auto self = FlatContCallArgs(op->args[0], keys);
                if (!keys.empty()) {
                    runtime::Array<BaseExpr> reverse_keys(keys.rbegin(), keys.rend());
                    runtime::Array<BaseExpr> call_args;
                    call_args.push_back(self);
                    call_args.push_back(InitializerList(std::move(reverse_keys), op->span));
                    call_args.push_back(op->args[2]);
                    return Call(op->checked_type(),
                                builtin::object___fused_setitem__(),
                                call_args,
                                op->span,
                                op->type_args);
                }
            } else if (op->op.same_as(builtin::object___getitem__())) {
                MXCHECK(op->args.size() == 2) << "internal error";
                runtime::Array<BaseExpr> keys;
                keys.push_back(op->args[1]);
                auto self = FlatContCallArgs(op->args[0], keys);
                if (!keys.empty()) {
                    runtime::Array<BaseExpr> reverse_keys(keys.rbegin(), keys.rend());
                    runtime::Array<BaseExpr> call_args;
                    call_args.push_back(self);
                    call_args.push_back(InitializerList(std::move(reverse_keys), op->span));
                    return Call(op->checked_type(),
                                builtin::object___fused_getitem__(),
                                call_args,
                                op->span,
                                op->type_args);
                }
            }
            return ExprMutator::VisitExpr_(op);
        }

    protected:
        static BaseExpr FlatContCallArgs(const BaseExpr &op, runtime::Array <BaseExpr> &keys) {
            auto *call_node = op.as<CallNode>();
            if (!call_node) {
                return op;
            }
            if (!call_node->op.same_as(builtin::object___getitem__())) {
                return op;
            }
            MXCHECK(call_node->args.size() == 2) << "internal error";
            keys.push_back(call_node->args[1]);
            return FlatContCallArgs(call_node->args[0], keys);
        }
    };

}  // namespace tbir::ir
